# -*- coding: utf-8 -*-
# created on 2016/12/22

from sympy import sympify, Eq, Add, solve, Mul, fraction, Abs, Poly, simplify, degree
from sympy.core.function import UndefinedFunction
from mathsolver.functions.base import BaseFunction, BaseFuncEq, BaseNumber, new_latex
from mathsolver.functions.sympy_utils import get_all_child, default_symbol


def zhouqi_dingyishi(zt):
    """hs060.函数周期定义式: f(x) 周期为 T"""
    relation = Eq(sympify('f(x)'), sympify('f(x+n*%s)' % zt))
    step = "由函数周期定义知 %s" % new_latex(relation)
    return relation, step


def zhouqi_dingyishi_general(fx, zt):
    """hs060.函数周期定义式: f(ax+b) 周期为 T"""
    sym = default_symbol(fx)
    fx_dingyi = fx.subs(sym, sym + sympify('n') * zt)
    res = Eq(fx, fx_dingyi)
    step = "根据函数周期性可知 %s" % new_latex(res)
    return res, step


def has_f(expr):
    """expr 是否是 f(*) 的格式"""
    return isinstance(type(expr), UndefinedFunction)


def has_minusf(expr):
    """expr 是否是 -f(*) 的格式"""
    return has_f(-expr)


def has_f_or_minusf(expr):
    """expr 是否是 f(*) 或者 -f(*) 的格式"""
    return has_f(expr) or has_f(-expr)


def is_fx(expr):
    """expr 是否是 f(x) 的格式"""
    return expr == sympify('f(x)')


def is_negfxplusb(expr):
    """expr 是否是 -f(x) + b 的格式，如果是则返回 b, b 可以为零"""
    symbol = default_symbol(expr)
    # 如果 b = 0
    if expr == sympify('-f(x)'):
        return 0
    b, negfx = expr.as_independent(symbol)
    if not b.has(symbol) and negfx.func is Mul and negfx.args[0] == -1 and is_fx(negfx.args[1]):
        return b


def is_bdivfx(expr):
    """expr 是否是 b/f(x) 的格式，如果是则返回 b"""
    b, fx = fraction(expr)
    symbol = default_symbol(expr)
    if not b.has(symbol) and is_fx(fx):
        return b


def is_fxplus_a(expr):
    """expr 是否是 f(x)+a 的格式"""
    symbol = default_symbol(expr)
    if expr.func is Add:
        a, fx = expr.args
        if not a.has(symbol) and is_fx(fx):
            return a


def is_fxplus_1(expr):
    """expr 是否是 f(x)+1 的格式"""
    return is_fxplus_a(expr) == 1


def is_fxminus_1(expr):
    """expr 是否是 f(x)-1 的格式"""
    return expr == sympify('f(x) -1')


def is_1minusfx(expr):
    """expr 是否是 1 - f(x) 的格式"""
    return expr == sympify('1 - f(x)')


def is_fxplus1divfxminus1(expr):
    """expr 是否是 (f(x)+1)/(f(x)-1) 的格式"""
    return simplify(sympify('(f(x)+1)/(f(x)-1)') - expr) == 0


def is_fxminus1divfxplus1(expr):
    """expr 是否是 (f(x)-1)/(f(x)+1) 的格式"""
    return simplify(sympify('(f(x)-1)/(f(x)+1)') - expr) == 0


def is_fxplus1div1minusfx(expr):
    """expr 是否是 (1+f(x))/(1-f(x)) 的格式"""
    return simplify(sympify('(1+f(x))/(1-f(x))') - expr) == 0


def is_1minusfxdivfxplus1(expr):
    """expr 是否是 (1-f(x))/(1+f(x)) 的格式"""
    return simplify(sympify('(1-f(x))/(1+f(x))') - expr) == 0


def is_fxplusa(expr):
    """expr 是否是 f(x+a) a != 0的格式，如果是则返回 a """
    if has_f(expr):
        symbol = default_symbol(expr)
        inner_f = expr.args[0]
        if inner_f.func is Add:
            a, x = inner_f.as_independent(symbol)
            if x == symbol and not a.has(symbol):
                return a


def is_faminusx(expr):
    """expr 是否是 f(a-x) 的格式，如果是则返回 a, a 可以为零 """
    if has_f(expr):
        symbol = default_symbol(expr)
        inner_f = expr.args[0]
        # 如果 a = 0
        if inner_f == -symbol:
            return 0
        if inner_f.func is Add:
            a, x = inner_f.as_independent(symbol)
            if x == -symbol and not a.has(symbol):
                return a


def is_m_minus_faminusx(expr):
    """expr 是否是 m - f(a-x) 的格式，如果是则返回 a; m,a 可以为零"""
    # m 为零
    if is_minus_faminusx(expr) is not None:
        return is_minus_faminusx(expr)

    symbol = default_symbol(expr)
    m, minus_faminusx = expr.as_independent(symbol)
    return is_minus_faminusx(minus_faminusx)


def is_fxplusbplusfx(expr):
    """expr 是否是 f(x+b) + f(x) 的格式，如果是则返回 b"""
    if expr.func is Add:
        fx, fxplusb = expr.args
        if is_fx(fx) and is_fxplusa(fxplusb) is not None:
            return is_fxplusa(fxplusb)


def is_minus_faminusx(expr):
    """expr 是否是 - f(a-x) 的格式，如果是则返回 a"""
    return is_faminusx(-expr)


def is_faxplusb(expr):
    """expr 是否是 f(a*x+b) 的格式，a != 0, b != 0, 如果是则返回 a, b"""
    if has_f(expr):
        inner_f = expr.args[0]
        if inner_f.func is Add:
            symbol = default_symbol(expr)
            b, ax = inner_f.as_independent(symbol)
            a, x = ax.as_independent(symbol)
            if not b.has(symbol) and x == symbol and not a.has(symbol):
                return a, b


def is_axplusb_anyb(expr):
    """expr 是否是 a*x+b 的格式，a != 0, b 可以为 0, 如果是则返回 a, b"""
    symbol = default_symbol(expr)
    try:
        if expr.is_polynomial(symbol) and degree(expr, gen=symbol) == 1:
            return Poly(expr, symbol).all_coeffs()
    except Exception:
        pass


def is_faxplusb_anyb(expr):
    """expr 是否是 f(a*x+b) 的格式，a != 0, b 可以为 0, 如果是则返回 a, b"""
    if has_f(expr):
        inner_f = expr.args[0]
        return is_axplusb_anyb(inner_f)


def hanshu_zhouqi_wrapper(eq):
    """求函数的周期, 输入 Eq"""
    left, right = eq
    zhouqi_solver = HanShuZhouQi().solver(BaseFuncEq([str(left), str(right)]))
    zt = zhouqi_solver.output[0].value
    steps = zhouqi_solver.steps
    return zt, steps


class HanShuZhouQi(BaseFunction):
    """hs061.求函数的周期"""

    def solver(self, *args):
        left, right = sympify(args[0].value)
        expr = left - right

        # 抽取出 f(x+a) a != 0
        f_xplusa_list = get_all_child(expr, is_fxplusa)

        if f_xplusa_list:
            f_xplusa = f_xplusa_list[0]
            a = is_fxplusa(f_xplusa)
            rest = solve(expr, f_xplusa)[0]
            # print rest
            if is_fx(rest):
                zt = Abs(a)
            elif is_fxplusa(rest) is not None:
                zt = Abs(a - is_fxplusa(rest))
            elif is_negfxplusb(rest) is not None or is_bdivfx(rest) is not None \
                    or is_fxplus1divfxminus1(rest) or is_1minusfxdivfxplus1(rest):
                zt = 2 * Abs(a)
            elif is_fxminus1divfxplus1(rest) or is_fxplus1div1minusfx(rest):
                zt = 4 * Abs(a)
            elif is_fxplusbplusfx(rest) is not None:
                b = is_fxplusbplusfx(rest)
                if b == 2 * a:
                    zt = 6 * Abs(a)
            else:
                raise ValueError('unknown type')
            self.output.append(BaseNumber(zt))
        else:
            # 如果不存在 f(x+a)，则抽取出 f(a*x+b),  a != 0, b != 0
            f_axplusb_list = get_all_child(expr, is_faxplusb)
            if f_axplusb_list:
                f_axplusb = f_axplusb_list[0]
                axplusb = f_axplusb.args[0]
                a, b = Poly(axplusb, default_symbol(expr)).all_coeffs()
                # TODO: f(ax) 的周期为 Abs(b/a)
                zt = Abs(b)
                t_fax = Abs(b / a)
            else:
                raise ValueError('未知表达式类型 %s = %s' % (new_latex(left), new_latex(right)))
            self.output.extend([BaseNumber(zt), BaseNumber(t_fax)])

        self.label.add("求函数的周期")
        self.steps.append(["", "由表达式 %s 可以计算出 f(x) 的最小正周期为 T = %s" % (args[0].printing(), new_latex(zt))])
        return self


if __name__ == '__main__':
    pass
